"""
a10_incremental_training.py - 音频分类模型增量训练脚本

功能：
1. 加载预训练模型和已有类别编码器
2. 添加新类别输出层并初始化权重
3. 冻结基础特征提取层，仅训练新分类头
4. 支持新旧数据混合训练防止灾难性遗忘
5. 保存增量训练后的模型和更新后的标签编码器
"""

import torch
import numpy as np
import pandas as pd
import yaml
from pathlib import Path
from typing import List, Dict, Any, Tuple

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
import torch.cuda.amp as amp
import copy

# 复用原有数据集类
from a04_model_training import AudioDataset, load_metadata, create_data_loaders
from utils.audio_cnn import AudioCNN


class IncrementalAudioDataset(Dataset):
    """支持新旧数据混合的增量数据集"""

    def __init__(self, old_data: Tuple[List[str], List[int]],
                 new_data: Tuple[List[str], List[int]],
                 data_dir: str,
                 processed_dir: str):
        """
        参数:
            old_data: (旧数据文件列表, 旧标签列表)
            new_data: (新数据文件列表, 新标签列表)
            data_dir: 数据根目录
            processed_dir: 预处理特征目录
        """
        self.old_files, self.old_labels = old_data
        self.new_files, self.new_labels = new_data
        self.data_dir = Path(data_dir)
        self.processed_dir = Path(processed_dir)

    def __len__(self):
        return len(self.old_files) + len(self.new_files)

    def __getitem__(self, idx):
        if idx < len(self.old_files):
            file, label = self.old_files[idx], self.old_labels[idx]
        else:
            file, label = self.new_files[idx - len(self.old_files)], self.new_labels[idx - len(self.old_files)]

        data = np.load(self.processed_dir / file)
        features = data['mel'] if isinstance(data, np.lib.npyio.NpzFile) else data
        return torch.FloatTensor(features).unsqueeze(0), torch.tensor(label, dtype=torch.long)


def load_incremental_data(cfg: Dict[str, Any],
                          new_class_names: List[str]) -> Tuple[DataLoader, DataLoader, LabelEncoder]:
    """
    加载增量训练数据
    返回:
        train_loader, val_loader, 更新后的标签编码器
    """
    # 加载原始数据
    old_files, old_labels, old_classes = load_metadata(cfg)

    # 加载新数据metadata (假设新数据有单独的metadata_new.csv)
    new_meta_path = Path(cfg['data_dir']) / 'metadata_new.csv'
    new_df = pd.read_csv(new_meta_path)

    # 合并标签编码器
    label_encoder = torch.load(Path(cfg['model_save_path']).parent / 'label_encoder.pth')
    new_classes = list(label_encoder.classes_) + new_class_names
    label_encoder = LabelEncoder().fit(new_classes)

    # 编码新标签 (原始类别自动沿用旧编码)
    new_df['label_encoded'] = label_encoder.transform(new_df['label'])
    new_files = new_df['file_path'].str.replace(r'^data[\\/]processed[\\/]', '', regex=True).tolist()
    new_labels = new_df['label_encoded'].tolist()

    # 创建混合数据集 (保留20%旧数据)
    old_train, _, old_train_labels, _ = train_test_split(
        old_files, old_labels, test_size=0.8, stratify=old_labels, random_state=42
    )

    # 划分新数据
    new_train, new_val, new_train_labels, new_val_labels = train_test_split(
        new_files, new_labels, test_size=cfg['val_ratio'], stratify=new_labels, random_state=42
    )

    # 创建混合数据集
    train_dataset = IncrementalAudioDataset(
        old_data=(old_train, old_train_labels),
        new_data=(new_train, new_train_labels),
        data_dir=cfg['data_dir'],
        processed_dir=cfg['processed_dir']
    )

    val_dataset = IncrementalAudioDataset(
        old_data=([], []),  # 验证集只用新数据
        new_data=(new_val, new_val_labels),
        data_dir=cfg['data_dir'],
        processed_dir=cfg['processed_dir']
    )

    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=cfg['batch_size'],
        shuffle=True,
        num_workers=cfg['num_workers']
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=cfg['batch_size'],
        num_workers=cfg['num_workers']
    )

    return train_loader, val_loader, label_encoder


class IncrementalAudioCNN(nn.Module):
    """支持增量学习的音频CNN"""

    def __init__(self, base_model: nn.Module, new_num_classes: int):
        super().__init__()
        # 复制基础模型结构
        self.feature_extractor = copy.deepcopy(base_model.features)

        # 替换分类头
        in_features = base_model.classifier[-1].in_features
        self.classifier = copy.deepcopy(base_model.classifier[:-1])  # 保留除最后一层外的结构
        self.classifier.append(nn.Linear(in_features, new_num_classes))

        # 冻结特征提取层
        for param in self.feature_extractor.parameters():
            param.requires_grad = False

        # 分类头初始化
        nn.init.kaiming_normal_(self.classifier[-1].weight)

    def forward(self, x):
        x = self.feature_extractor(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)


def incremental_train(cfg: Dict[str, Any], new_class_names: List[str]):
    """执行增量训练"""
    # 设备设置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 加载基础模型
    base_model = AudioCNN(
        num_classes=len(torch.load(Path(cfg['model_save_path']).parent / 'label_encoder.pth').classes_))
    base_model.load_state_dict(torch.load(cfg['model_save_path']))
    base_model.to(device)

    # 准备增量模型
    model = IncrementalAudioCNN(
        base_model=base_model,
        new_num_classes=len(new_class_names) + len(base_model.classifier[-1].weight)
    ).to(device)

    # 加载数据
    train_loader, val_loader, label_encoder = load_incremental_data(cfg, new_class_names)

    # 训练配置
    optimizer = optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),  # 只训练解冻参数
        lr=cfg['lr'] * 0.1,  # 使用更小的学习率
        weight_decay=cfg['weight_decay']
    )
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg['epochs'])
    criterion = nn.CrossEntropyLoss()
    scaler = amp.GradScaler(enabled=cfg['use_amp'])

    # 训练循环
    best_acc = 0.0
    for epoch in range(1, cfg['epochs'] + 1):
        model.train()
        total_loss = 0.0

        for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch}/{cfg['epochs']}"):
            inputs, labels = inputs.to(device), labels.to(device)

            with amp.autocast(enabled=cfg['use_amp']):
                outputs = model(inputs)
                loss = criterion(outputs, labels)

            optimizer.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            total_loss += loss.item()

        # 验证
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        val_acc = 100 * correct / total
        print(f"Epoch {epoch} - Loss: {total_loss / len(train_loader):.4f} | Val Acc: {val_acc:.2f}%")

        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), cfg['model_save_path'].replace('.pth', '_inc.pth'))
            torch.save(label_encoder, Path(cfg['model_save_path']).parent / 'label_encoder_inc.pth')

    print(f"增量训练完成! 最佳验证准确率: {best_acc:.2f}%")


if __name__ == "__main__":
    # 加载配置
    with open("configs/default.yaml", "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    # 新类别定义 (根据实际需求修改)
    new_classes = ["new_class1", "new_class2"]

    # 执行增量训练
    incremental_train(cfg, new_classes)